6d7422
@@ -760,29 +760,21 @@
SSLConnectionSocketFactory getTwoWaySSLSocketFactory() throws SQLException {
   }
 
   // Lookup the delegation token. First in the connection URL, then Configuration
-  private String getClientDelegationToken(Map<String, String> jdbcConnConf)
-      throws SQLException {
+  private String getClientDelegationToken(Map<String, String> jdbcConnConf) throws SQLException {
     String tokenStr = null;
-    if (JdbcConnectionParams.AUTH_TOKEN.equalsIgnoreCase(jdbcConnConf.get(JdbcConnectionParams.AUTH_TYPE))) {
-      // check delegation token in job conf if any
+    if (!JdbcConnectionParams.AUTH_TOKEN.equalsIgnoreCase(jdbcConnConf.get(JdbcConnectionParams.AUTH_TYPE))) {
+      return null;
+    }
+    DelegationTokenFetcher fetcher = new DelegationTokenFetcher();
+    try {
+      tokenStr = fetcher.getTokenStringFromFile();
+    } catch (IOException e) {
+      LOG.warn("Cannot get token from environment variable $HADOOP_TOKEN_FILE_LOCATION=" +
+              System.getenv(UserGroupInformation.HADOOP_TOKEN_FILE_LOCATION));
+    }
+    if (tokenStr == null) {
       try {
-        if (System.getenv(UserGroupInformation.HADOOP_TOKEN_FILE_LOCATION) != null) {
-          try {
-            Credentials cred = new Credentials();
-            DataInputStream dis = new DataInputStream(new FileInputStream(System.getenv(UserGroupInformation
-                    .HADOOP_TOKEN_FILE_LOCATION)));
-            cred.readTokenStorageStream(dis);
-            dis.close();
-            Token<? extends TokenIdentifier> token = cred.getToken(new Text("hive"));
-            tokenStr = token.encodeToUrlString();
-          } catch (IOException e) {
-            LOG.warn("Cannot get token from environment variable $HADOOP_TOKEN_FILE_LOCATION=" +
-                    System.getenv(UserGroupInformation.HADOOP_TOKEN_FILE_LOCATION));
-          }
-        }
-        if (tokenStr == null) {
-          tokenStr = SessionUtils.getTokenStrForm(HiveAuthConstants.HS2_CLIENT_TOKEN);
-        }
+        return fetcher.getTokenFromSession();
       } catch (IOException e) {
         throw new SQLException("Error reading token ", e);
       }
@@ -790,6 +782,34 @@
private String getClientDelegationToken(Map<String, String> jdbcConnConf)
     return tokenStr;
   }
 
+  static class DelegationTokenFetcher {
+    String getTokenStringFromFile() throws IOException {
+      if (System.getenv(UserGroupInformation.HADOOP_TOKEN_FILE_LOCATION) == null) {
+        return null;
+      }
+      Credentials cred = new Credentials();
+      try (DataInputStream dis = new DataInputStream(new FileInputStream(System.getenv(UserGroupInformation
+              .HADOOP_TOKEN_FILE_LOCATION)))) {
+        cred.readTokenStorageStream(dis);
+      }
+      return getTokenFromCredential(cred, "hive");
+    }
+
+    String getTokenFromCredential(Credentials cred, String key) throws IOException {
+      Token<? extends TokenIdentifier> token = cred.getToken(new Text(key));
+      if (token == null) {
+        LOG.warn("Delegation token with key: [hive] cannot be found.");
+        return null;
+      }
+      return token.encodeToUrlString();
+    }
+
+    String getTokenFromSession() throws IOException {
+      LOG.debug("Fetching delegation token from session.");
+      return SessionUtils.getTokenStrForm(HiveAuthConstants.HS2_CLIENT_TOKEN);
+    }
+  }
+
   private void openSession() throws SQLException {
     TOpenSessionReq openReq = new TOpenSessionReq();
 
